;; Tests of the matrix chain multiplication functions
;; Copyright (C) 2022  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.matrix-chain-test
  (:require
   [clojure.test :refer [deftest is testing]]
   [mazama.linalg.matrix-chain :refer [matrix-chain multiply]]))

(defn- fncols [A]
  (:ncols A))

(defn- fnrows [A]
  (/ (count (:data A)) (:ncols A)))

(defn- make-matrix [nrows ncols]
  {:ncols ncols :data (repeatedly (* nrows ncols) #(- (rand-int 100) 50))})

(defn- make-matrix-sequence [dims]
  (map make-matrix dims (rest dims)))

(defn columns
  "Return a sequence of columns, each column being a sequence of elements."
  ([{:keys [ncols data]}]
   (columns ncols data))
  ([ncols data]
   (letfn [(f [i ncols data]
             (if (< i ncols)
               (lazy-seq (cons (take-nth ncols data) (f (inc i) ncols (rest data))))
               []))]
     (f 0 ncols (seq data)))))

(defn rows
  "Return a sequence of rows, each row being a sequence."
  ([{:keys [ncols data]}]
   (rows ncols data))
  ([ncols data]
   (partition ncols data)))

(defn dot-product
  [u v]
  (reduce + 0 (mapv * u v)))

(defn- f* [A B]
  {:ncols (:ncols B)
   :data (for [row (rows A) col (columns B)] (dot-product row col))})

(deftest small-test
  (testing "nil"
    (let [chain (matrix-chain fnrows fncols nil)]
      (is (empty? chain))))
  (testing "empty"
    (let [chain (matrix-chain fnrows fncols ())]
      (is (empty? chain))))
  (testing "one"
    (let [A (make-matrix 4 3)
          coll (vector A)
          chain (matrix-chain fnrows fncols coll)]
      (is (= A chain))))
  (testing "two"
    (let [A (make-matrix 7 6)
          B (make-matrix 6 3)
          coll (vector A B)
          chain (matrix-chain fnrows fncols coll)]
      (is (= coll chain)))))

(deftest wikipedia-test
  (testing "https://en.wikipedia.org/wiki/Matrix_chain_multiplication"
    (let [A (make-matrix 10 30)
          B (make-matrix 30 5)
          C (make-matrix 5 60)
          coll (vector A B C)
          chain (matrix-chain fnrows fncols coll)]
      (is (= [[A B] C] chain)))))

(deftest mcm-test
  (testing "https://www.mimuw.edu.pl/~erykk/algovis/mcm.html"
    (let [dims [3 8 1 2 5 7 6 1 8 4 7 5 2 1 9]
          [A B C D E F G H I J K L M N :as coll] (make-matrix-sequence dims)
          chain (matrix-chain fnrows fncols coll)]
      (is (= [[[A B][[C[D[E[F G]]]][[[[[H I]J]K]L]M]]]N] chain)))))

(deftest multiply-test
  (testing "https://en.wikipedia.org/wiki/Matrix_multiplication#Resource_allocation_in_economics"
    (let [A {:ncols 3 :data [1 0 1 2 1 1 0 1 1 1 1 2]}
          B {:ncols 3 :data [1 2 1 2 3 1 4 2 2]}
          C {:ncols 1 :data [100 80 60]}
          chain (matrix-chain fnrows fncols (list A B C))]
      (is (= [A [B C]] chain))
      (is (= {:ncols 1 :data [1000 1820 1180 2180]} (multiply f* chain))))))
